// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <arrow/array.h>
#include <arrow/c/abi.h>
#include <arrow/c/bridge.h>
#include <arrow/python/pyarrow.h>
#include <arrow/record_batch.h>

#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include <pollux/python/init/PyInit.h>
#include <pollux/python/vector/PyVector.h>
#include <pollux/vector/arrow/bridge.h>

namespace py = pybind11;

/// This module adds two functions `to_pollux()` and `to_arrow()` that allow the
/// conversion between Pollux Vectors and Arrow Arrays from a Python program. It
/// works by extracting the Arrow C structures from the Arrow C++ Array, then
/// using Pollux's Arrow bridge to convert it to a Pollux Vector (and vice-versa).
PYBIND11_MODULE(arrow, m) {
  using namespace kumo;

  py::module::import("pypollux.vector");

  arrow::py::import_pyarrow();
  pollux::py::initializePolluxMemory();

  static auto rootPool = pollux::memory::memoryManager()->addRootPool();
  static auto leafPool = rootPool->addLeafChild("py_pollux_arrow_pool");

  /// Converts a pyarrow.Array into a pypollux.vector.Vector using Pollux's arrow
  /// bridge.
  m.def("to_pollux", [](py::object& batchObject) {
    ArrowSchema schema;
    ArrowArray data;
    std::shared_ptr<arrow::Array> array;
    std::shared_ptr<arrow::RecordBatch> batch;

    if (arrow::py::is_array(batchObject.ptr())) {
      array = *arrow::py::unwrap_array(batchObject.ptr());
      auto statusType = arrow::ExportType(*array->type(), &schema);
      auto statusArray = arrow::ExportArray(*array, &data);

      if (!statusArray.ok() || !statusType.ok()) {
        throw std::runtime_error("Unable to convert Arrow Array to C ABI.");
      }
    } else if (arrow::py::is_batch(batchObject.ptr())) {
      batch = *arrow::py::unwrap_batch(batchObject.ptr());
      auto statusType = arrow::ExportSchema(*batch->schema(), &schema);
      auto statusArray = arrow::ExportRecordBatch(*batch, &data);

      if (!statusArray.ok() || !statusType.ok()) {
        throw std::runtime_error(
            "Unable to convert Arrow RecorBatch to C ABI.");
      }
    } else {
      throw std::runtime_error("Unknown input Arrow structure.");
    }
    return pollux::py::PyVector{
        pollux::import_from_arrow_as_viewer(schema, data, leafPool.get()), leafPool};
  });

  /// Converts a pypollux.vector.Vector to a pyarrow.Array using Pollux's arrow
  /// bridge.
  m.def("to_arrow", [](pollux::py::PyVector& vector) {
    ArrowSchema schema;
    ArrowArray data;

    pollux::export_to_arrow(vector.vector(), schema);
    pollux::export_to_arrow(vector.vector(), data, leafPool.get());

    auto arrowType = *arrow::ImportType(&schema);
    auto arrowArray = *arrow::ImportArray(&data, arrowType);
    return py::reinterpret_steal<py::object>(arrow::py::wrap_array(arrowArray));
  });
}
